from math import sqrt
import sys, os
sys.path.append(os.path.dirname(os.path.dirname(os.path.dirname(os.path.abspath(__file__)))))
import numpy as np
import cv2, json, pandas as pd, seaborn as sns
import matplotlib.pyplot as plt

from protocol.decorators import edp_service, instant, instant_req
# @edp_service
class FeatureAnalysis:
    name = "feature analysis"
    desc = ''
    category = "feature"

    tasks = [
        'collect_stage2_feat: (List<Pose>, List<(Int, Int, Bool)>, List<(Int, Int, Bool)>, Frequency) -> (List<(Float, Float)>, List<(Float, Float)>)',
        'collect_state2_feat_all: (Str, Map<Str, Int>)->(List<List<(Float, Float)>>, List<List<(Float, Float)>>)',
        'feat2_clean: Feat2 -> Feat2',
        'feat2_table: (Feat2, Str)->Str',
        'vis_feat2: (Feat2, Str)->Void',
        'vis_feat3: (Feat3, Str)->Void',
    ]
    @instant
    def collect_stage2_feat(self, inputs):
        pose, sl, sr, freq = inputs
        freq = freq // 10 - 1
        res = []
        for s, e, _ in sl:
            x1,y1,a1 = pose[s][0]
            x2,y2,a2 = pose[e][0]
            diff = sqrt((x2-x1)**2 + (y2-y1)**2)
            da = a2 - a1
            res.append([False, freq, diff, da])
        for s, e, _ in sr:
            x1,y1,a1 = pose[s][0]
            x2,y2,a2 = pose[e][0]
            diff = sqrt((x2-x1)**2 + (y2-y1)**2)
            da = a2 - a1
            res.append([True, freq, diff, da])
        return res
    @instant
    def collect_state2_feat_all(self, inputs):
        dd = inputs[0]
        freqdc = inputs[1]
        res = []
        for vd in os.listdir(dd):
            if vd not in freqdc:
                print(f"no freq: {vd}")
                continue
            freq = freqdc[vd]
            vd = dd + "/" + vd
            if not os.path.exists(vd+"/pose.json"):
                print(f"no pose: {vd}")
                continue
            pose = json.load(open(vd+"/pose.json","r"))
            if not os.path.exists(vd+"/lstim_seg.json"):
                print(f"no lstim_seg: {vd}")
                continue
            lstim = json.load(open(vd+"/lstim_seg.json","r"))
            if not os.path.exists(vd+"/rstim_seg.json"):
                print(f"no rstim_seg: {vd}")
                continue
            rstim = json.load(open(vd+"/rstim_seg.json","r"))
            res.append(self.collect_stage2_feat((pose, lstim, rstim, freq)))
        featall = [[[] for _ in range(6)] for _ in range(2)]
        for r in res:
            for each in r:
                if each[0]:
                    featall[1][each[1]].append(each[2:])
                else:
                    featall[0][each[1]].append(each[2:])
        return featall
    def feat2_clean(self, inputs):
        # 去除异常值
        featall = inputs
        for i in range(2):
            for j in range(6):
                if len(featall[i][j]) > 0:
                    mean_diff = np.mean([x[0] for x in featall[i][j]])
                    std_diff = np.std([x[0] for x in featall[i][j]])
                    mean_angle = np.mean([x[1] for x in featall[i][j]])
                    std_angle = np.std([x[1] for x in featall[i][j]])
                    featall[i][j] = [x for x in featall[i][j] if x[0] > mean_diff - std_diff and x[0] < mean_diff + std_diff and x[1] > mean_angle - std_angle and x[1] < mean_angle + std_angle]
        return featall
    def feat2_table(self, inputs):
        featall = inputs[0]
        csvd = inputs[1]
        os.makedirs(csvd, exist_ok=True)
        fpath = csvd + '/feat2.csv'
        df = pd.DataFrame(columns=['Position', 'Frequency', 'Mean Diff', 'Std Diff', 'Mean Angle', 'Std Angle'])
        for i in range(2):
            for j in range(6):
                if len(featall[i][j]) > 0:
                    mean_diff = np.mean([x[0] for x in featall[i][j]])
                    mean_angle = np.mean([x[1] for x in featall[i][j]]) 
                    std_diff = np.std([x[0] for x in featall[i][j]])
                    std_angle = np.std([x[1] for x in featall[i][j]])
                    df.loc[len(df)] = {'Position': i, 'Frequency': j*10+10, 'Mean Diff': mean_diff, 'Std Diff': std_diff, 'Mean Angle': mean_angle, 'Std Angle': std_angle}
                else:
                    df.loc[len(df)] = {'Position': i, 'Frequency': j*10+10, 'Mean Diff': 0, 'Std Diff': 0, 'Mean Angle': 0, 'Std Angle': 0}
        df.to_csv(fpath, index=False)
        print(f'save table to {fpath}')
        return fpath
    def vis_feat2(self, inputs):
        featall = inputs[0]
        visd = inputs[1]
        os.makedirs(visd, exist_ok=True)
        freqs = list(range(10, 70, 10))

        # 设置 seaborn 样式
        sns.set_style("whitegrid")
        sns.set_context("notebook", font_scale=1.5)
        plt.rcParams['legend.fontsize'] = 14
        plt.rcParams['axes.labelsize'] = 16
        plt.rcParams['axes.titlesize'] = 18

        for i in range(2):
            lr = 'R' if i else 'L'
            for j in range(6):
                fname = f'diff_{j * 10 + 10}{lr}.png'
                fpath = visd + '/' + fname
                plt.figure(figsize=(8, 6))
                data = featall[i][j]
                if len(data) > 0:
                    diffs = [x[0] for x in data]
                    df = pd.DataFrame({
                        'sample': range(len(diffs)),
                        'diff': diffs
                    })
                    sns.scatterplot(data=df, x='sample', y='diff', alpha=0.5)
                    plt.axhline(y=np.mean(diffs), color='r', linestyle='--', label='mean')
                    plt.legend()
                plt.title(f'{j*10+10}Hz {lr}')
                plt.xlabel('sample')
                plt.ylabel('diff')
                plt.savefig(fpath)
                plt.close()

                fname = f'angle_{j * 10 + 10}{lr}.png'
                fpath = visd + '/' + fname
                plt.figure(figsize=(8, 6))
                data = featall[i][j]
                if len(data) > 0:
                    angles = [x[1] for x in data]
                    df = pd.DataFrame({
                        'sample': range(len(angles)), 
                        'angle': angles
                    })
                    sns.scatterplot(data=df, x='sample', y='angle', alpha=0.5)
                    plt.axhline(y=np.mean(angles), color='r', linestyle='--', label='mean')
                    plt.legend()
                plt.title(f'{j*10+10}Hz {lr}')
                plt.xlabel('sample')
                plt.ylabel('angle')
                plt.savefig(fpath)
                plt.close()

            # 修改为箱线图
            fname = f'overall_diff_{lr}.png'
            fpath = visd + '/' + fname
            plt.figure(figsize=(12, 6))
            diff_data = []
            for j in range(6):
                data = featall[i][j]
                if len(data) > 0:
                    diffs = [x[0] for x in data]
                    for diff in diffs:
                        diff_data.append({
                            'Frequency': j * 10 + 10,
                            'Diff': diff
                        })
            df_diff = pd.DataFrame(diff_data)
            if not df_diff.empty:
                sns.boxplot(data=df_diff, x='Frequency', y='Diff', color='skyblue')
                # 添加均值点
                sns.pointplot(data=df_diff, x='Frequency', y='Diff', color='red', markers='D', scale=0.7)
            plt.xlabel('Frequency (Hz)')
            plt.ylabel('Diff')
            plt.title(f'Overall Diff Statistics for {lr} Stimulation')
            plt.savefig(fpath, bbox_inches='tight', dpi=300)
            plt.close()

            fname = f'overall_angle_{lr}.png'
            fpath = visd + '/' + fname
            plt.figure(figsize=(12, 6))
            angle_data = []
            for j in range(6):
                data = featall[i][j]
                if len(data) > 0:
                    angles = [x[1] for x in data]
                    for angle in angles:
                        angle_data.append({
                            'Frequency': j * 10 + 10,
                            'Angle': angle
                        })
            df_angle = pd.DataFrame(angle_data)
            if not df_angle.empty:
                sns.boxplot(data=df_angle, x='Frequency', y='Angle', color='lightgreen')
                # 添加均值点
                sns.pointplot(data=df_angle, x='Frequency', y='Angle', color='red', markers='D', scale=0.7)
            plt.xlabel('Frequency (Hz)')
            plt.ylabel('Angle')
            plt.title(f'Overall Angle Statistics for {lr} Stimulation')
            plt.savefig(fpath, bbox_inches='tight', dpi=300)
            plt.close()

        # 修改为箱线图
        plt.figure(figsize=(15, 6))    
        diff_data = []
        for j in range(6):
            # Left stimulation
            data = featall[0][j]
            if len(data) > 0:
                diffs = [x[0] for x in data]
                for diff in diffs:
                    diff_data.append({
                        'Frequency': j * 10 + 10,
                        'Diff': diff,
                        'Stimulation': 'L'
                    })
            # Right stimulation
            data = featall[1][j]
            if len(data) > 0:
                diffs = [x[0] for x in data]
                for diff in diffs:
                    diff_data.append({
                        'Frequency': j * 10 + 10,
                        'Diff': diff,
                        'Stimulation': 'R'
                    })
        df_diff_compare = pd.DataFrame(diff_data)
        if not df_diff_compare.empty:
            sns.boxplot(data=df_diff_compare, x='Frequency', y='Diff', hue='Stimulation',
                       palette={'L': 'skyblue', 'R': 'lightblue'})
            # 添加均值点
            sns.pointplot(data=df_diff_compare, x='Frequency', y='Diff', hue='Stimulation',
                         palette={'L': 'red', 'R': 'darkred'}, markers='D', scale=0.7)
        plt.title('Left vs Right Stimulation - Diff Comparison')
        plt.savefig(visd + '/lr_diff.png', bbox_inches='tight', dpi=300)
        plt.close()

        plt.figure(figsize=(15, 6))
        angle_data = []
        for j in range(6):
            # Left stimulation
            data = featall[0][j]
            if len(data) > 0:
                angles = [x[1] for x in data]
                for angle in angles:
                    angle_data.append({
                        'Frequency': j * 10 + 10,
                        'Angle': angle,
                        'Stimulation': 'L'
                    })
            # Right stimulation
            data = featall[1][j]
            if len(data) > 0:
                angles = [x[1] for x in data]
                for angle in angles:
                    angle_data.append({
                        'Frequency': j * 10 + 10,
                        'Angle': angle,
                        'Stimulation': 'R'
                    })
        df_angle_compare = pd.DataFrame(angle_data)
        if not df_angle_compare.empty:
            sns.boxplot(data=df_angle_compare, x='Frequency', y='Angle', hue='Stimulation',
                       palette={'L': 'lightgreen', 'R': 'palegreen'})
            # 添加均值点
            sns.pointplot(data=df_angle_compare, x='Frequency', y='Angle', hue='Stimulation',
                         palette={'L': 'red', 'R': 'darkred'}, markers='D', scale=0.7)
        plt.title('Left vs Right Stimulation - Angle Comparison')
        plt.savefig(visd + '/lr_angle.png', bbox_inches='tight', dpi=300)
        plt.close()
        print(f'save figs to {visd}')

    def vis_feat3(self, inputs):
        ...

def cross(a1, a2):
    res = []
    for x,y in zip(a1,a2):
        res.append(x)
        res.append(y)
    return res


if __name__ == '__main__':
    o = FeatureAnalysis()
    prj = r"C:\Users\songy\Desktop\fyp-songy-2025-05-17"
    dd = r"C:\Users\songy\Desktop\data-0518"
    if not os.path.exists(dd + '/feat2.json'):
        freqdc = instant_req("http://127.0.0.1:5060" + "/task/config_get_freqs", prj)
        feat = o.collect_state2_feat_all((dd, freqdc))
        json.dump(feat, open(dd + '/feat2.json', 'w'))
    else:
        feat = json.load(open(dd + '/feat2.json', 'r'))
    feat = o.feat2_clean(feat)
    csvd = r"E:\25spring\FYP\pymodules\figs" + "/feat2"
    o.feat2_table((feat, csvd))
    visd = r"E:\25spring\FYP\pymodules\figs" + "/feat2"
    o.vis_feat2((feat, visd))

